import cv2
from paddleocr import PaddleOCR, draw_ocr
import paddle
from PIL import Image
import os
import time

# 初始化 PaddleOCR
start_time=time.time()
ocr = PaddleOCR(use_angle_cls=True, lang="ch")
print("模型加载时间:",time.time()-start_time)
print(f"Current device: {paddle.device.get_device()}")

# 图像文件路径
img_path = './images/general_ocr_rec_001.png'

# 检查文件是否存在
if not os.path.exists(img_path):
    print(f"Error: Image file {img_path} does not exist.")
else:
    try:
        # 读取图像
        img = cv2.imread(img_path)
        if img is None:
            print(f"Error: Failed to load image {img_path}.")
        else:
            # 进行 OCR 处理
            start_time=time.time()
            result = ocr.ocr(img, cls=True)
            print("模型推理时间:",time.time()-start_time)
            for idx in range(len(result)):
                res = result[idx]
                for line in res:
                    print(line)

            # 显示结果
            from PIL import Image
            result = result[0]
            image = Image.open(img_path).convert('RGB')
            boxes = [line[0] for line in result]
            txts = [line[1][0] for line in result]
            scores = [line[1][1] for line in result]
            im_show = draw_ocr(image, boxes, txts, scores, font_path='./doc/fonts/simfang.ttf')
            im_show = Image.fromarray(im_show)
            im_show.save('result.jpg')
    except Exception as e:
        print(f"Error: {e}")
